/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Definition of the Snapshot class. \ingroup sfrhist

// $Id$

#ifndef __snapshot__
#define __snapshot__

#include "config.h"
#include <vector>
#include <map>
#include <set>
#include <string>
#include "biniostream.h"
#include "message.h"
#include "mcrx-types.h"
#include "mcrx-units.h"
#include "particle_set.h"
#include "gadgetquant.h"
#include <cassert>
#include <exception>
#include "mcrx-debug.h"
#include <numeric>


class Preferences;
class History;
namespace CCfits {
  class Table;
};

namespace mcrx {
  class Snapshot;
};

namespace H5 {
  class CommonFG;
};

/// @{
/// \ingroup sfrhist


/** Contains the data in a snapshot from the N-body code.  At this
    time, only GADGET snapshots work. \ingroup sfrhist */
class mcrx::Snapshot : public message {
public:
  /// This defines the particle species.
  enum species { gas=0, halo=1, disk=2, bulge=3, nstar=4, bh=5 };
  static const int n_species=6;

  // The following are private because they are conceptually const and
  // should be accessed through the corresponding access functions.
private:
  /// Vector which keeps pointers to the sets, so they can be indexed
  /// using the enums. This should only be accessed through the sets()
  /// function.
  std::vector<particle_set_generic*> sets_;
  /// This set contains all species that are star_particle_sets,
  /// because some operations work on thos species. This should only
  /// be accessed through the star_species() function.
  std::set<species> star_species_;
  static std::vector<std::string> species_names_; ///< Vector of names of the species

protected:
  const Preferences& prefs_;
  const Preferences& user_prefs_;

  double time_; ///< Snapshot time.
  double a_; ///< Expansion factor (1 if not cosmological)
  std::string name_;

  bool byte_swap;

  /** \name Units
      These refer to what the units of the arrays are, in terms of our
      fiducial units kpc/M_sun/yr. */
  /// @{ 
  double massunit_;
  double lengthunit_;
  double timeunit_;
  double G_;
  /// The standard unit map. This is updated whenever any of the above change
  T_unit_map units_; 
  /// @}

  /** Arrays for gas particles.
      Quantities are stored in physical units, not simulation units.
      Particles are ordered in the arrays in order of increasing ID
      number, but these may not be a continuous sequence. */
  gas_particle_set gas_particles;

  /** Arrays for dark-matter halo particles. */
  dm_particle_set halo_particles;

  /** Arrays for disk (pre-merger) star particles. */
  star_particle_set disk_particles;

  /** Arrays for bulge (pre-merger) star particles. */
  star_particle_set bulge_particles;

  /** Arrays for new star particles. */
  star_particle_set nstar_particles;

  /** Arrays for black-hole particles. */
  bh_particle_set bh_particles;

  /** Handy utility function to read a field from the file into a
      vector, using the order vector to order the data. The type of the
      vector is deduced, but the type of the field must be supplied. */
  template <typename T_field, typename T_file>
  bool read_field (binifstream& stream, int n, 
		   std::vector<T_field>& v, std::vector<int>& read_order, 
		   T_file, bool check = true) const;
  /// Skip a field in the snapshot file, using the formatting length info.
  bool skip_field (binifstream&, int) const;
  /// Check that the field header is consistent with a record of n bytes.
  bool check_field_header(binifstream& stream, int n) const; 
  void setup_ordering(std::vector<int>& read_order,
		      particle_set_generic& set,
		      std::map<unsigned int, int> shuffle, int);
  void calculate_temperatures(std::vector<float>& u_ngas,
			      std::vector<float>& ne,
			      std::vector<float>& tp);
  template<typename S>
  void decode_scannapieco(const std::vector<blitz::TinyVector<float,12> >&,
			  S&);
  void assign_softening(species, const std::string&, std::vector<T_float>&);
  void init_disk_bulge_halo();
  
  template <typename T_field>
  bool read_dataset (const H5::CommonFG& f, const std::string& datasetname,
		     std::vector<T_field>& field,
		     const std::vector<int>& read_order=std::vector<int>()) const;
  template<typename T, int N>
  bool read_dataset (const H5::CommonFG& f, const std::string& datasetname,
		     std::vector<blitz::TinyVector<T, N> >& field,
		     const std::vector<int>& read_order) const;

  // Load functions are called from constructor
  bool load(const std::string& snapfn);
  bool loadgadget(const std::string& snapfn);
  bool loadgadget_hdf5(const std::string& snapfn);
  bool loadtipsy(const std::string& snapfn);
  bool loadtipsy_ascii(const std::string& snapfn);
  bool loadtipsy_bin(const std::string& snapfn);
  bool loadenzo(const std::string& snapfn);
  bool loadart(const std::string& snapfn);
  void check_parent_ID_integrity() const;

  /** Convert expansion factor to h^-1 time. Used when loading
      comoving snapshots. Thanks to Mark Vogelsberger for providing
      this. */
  static double a2time(double a, double OmegaM, double OmegaL, 
		       double UnitTime_in_s);
  
public:
  Snapshot(const Preferences&, const Preferences&, 
	   const std::string& file, const std::string& name);
  virtual ~Snapshot () {};


  bool makephysicalunits();

  bool isgadget(const std::string& snapfn);

  // This version is potentially unsafe because you can index the
  // vector with an integer...
  const std::vector<particle_set_generic*>& sets() const {return sets_;};
  particle_set_generic& set(species s) {return *sets()[s];};
  const particle_set_generic& set(species s) const {return *sets()[s];};
  const std::set<species>& star_species() const {return star_species_;};
  static std::string species_name(species s);

  // Returns vector index for a given ID and species.
  int find_ID(species s, unsigned int id) const {
    return sets()[s]->find_ID(id);
  };

  // Returns particle set and index for a given ID. This is the most
  // generic particle retrieval function.
  std::pair<species, int> find_ID(int id) const;

  /// Returns the number of particles of a certain species.
  int n(species s) const { return sets()[s]->n();};

  /// Returns the total number of particles
  int ntot() const { return n(gas)+n(disk)+n(bulge)+n(halo)+n(nstar)+n(bh);};

  double time() const { return time_; };
  const std::string& name() const { return name_; };
  bool  swap_byte_order() const {return byte_swap;};
  const T_unit_map& units () const {return units_;};

  /** Addition operator is used for concatenating the particles in two
      snapshots. This is used for the Arepo/Gadget multifile
      snapshots, so it assumes that the header info is identical in
      them and doesn't do anything with the header in the rhs. */
  Snapshot& operator+=(const Snapshot& rhs);
};

/// @}

/** Calculates total metallicity from the Scannapieco 12-vector of
    different element species. */
template<typename S>
void 
mcrx::Snapshot::decode_scannapieco(const std::vector<blitz::TinyVector<float,12> >& metals,
				   S& set)
{
  set.z_.resize(set.n());
  for(int i=0; i< set.n(); ++i) {
    const blitz::TinyVector<float,12>& zz(metals[i]);
    set.z_[i] = (zz[1]+zz[2]+zz[3]+zz[4]+zz[5]+zz[7]+
		 zz[8]+zz[9]+zz[10]+zz[11])/set.m_[i];
  }
  blitz::TinyVector<float,12> zero; zero=0.0;
  const blitz::TinyVector<float,12> tot=std::accumulate(metals.begin(),
							metals.end(),
							zero);
  DEBUG(1, std::cout << "Total mass of Scannapieco metals: " << tot << std::endl;);
}

#endif

